#include "mm.h"
#include "slab.h"
#include "mmu.h"
#include "sysregs.h"


static MemoryManagement mmm;
static pgd_t* new_pgdp;

static void early_memory_init() {
	mmm.bitmap_num = TOTAL_MEMORY / PAGE_SIZE / 64;
	mmm.bitmap_size = mmm.bitmap_num * 8;
	mmm.bitmap_p = (void*)(KERNEL_START + LOW_MEMORY);		

	for (u64 i = 0; i < mmm.bitmap_num; i++)
		mmm.bitmap_p[i] = 0;

	mmm.pages_num = TOTAL_MEMORY / PAGE_SIZE;
	mmm.pages_size = mmm.pages_num * sizeof(Page);
	mmm.pages_p = (Page*)((u64)mmm.bitmap_p + mmm.bitmap_size);

	for (u64 i = 0; i < mmm.pages_num; i ++) {
		mmm.pages_p[i].p_addr = (void*)(PAGE_SIZE * i);
		mmm.pages_p[i].count  = 0;
	}

	mmm.end_addr = (void*)((u64)PAGE_ALIGN(mmm.pages_p + mmm.pages_num + 1) + PAGE_SIZE); 

	u64 byte_used = (u64)mmm.end_addr - KERNEL_START;
	u64 pages_used = byte_used >> PAGE_SHIFT;
	for (u64 i = 0; i < pages_used; i++) {
		u64* bitmap_p = mmm.bitmap_p + (i / 64);
		*bitmap_p |= UL(1) << (i % (sizeof(u64)*8));
		Page* page_p = mmm.pages_p + i;
		page_p->count += 1;
	}
}


static void init_mapping() {
	new_pgdp = (pgd_t*)alloc_one_page();
	u64 start, end, size;

	start = (u64)text_begin;
	end   = (u64)rodata_begin;
	__create_pgd_mapping(new_pgdp, VirtToPhy(start), start, end-start, PAGE_KERNEL_ROX);

	start = (u64)rodata_begin;
	end   = (u64)data_begin;
	__create_pgd_mapping(new_pgdp, VirtToPhy(start), start, end-start, PAGE_KERNEL_RO);

	start = (u64)data_begin;
	size  = (u64)TOTAL_MEMORY - (u64)VirtToPhy(start);
	__create_pgd_mapping(new_pgdp, VirtToPhy(start), start, size, PAGE_KERNEL);

	start = (u64)(PBASE);
	__create_pgd_mapping(new_pgdp, start, PhyToVirt(start), DEVICE_SIZE, PROT_DEVICE_nGnRnE);
}


static void switch_pt() {

	write_sysreg(new_pgdp, ttbr1_el1);
	isb();

	asm("tlbi vmalle1");
	isb();

	asm("ic iallu");
	dsb(nsh);
	isb();
}


void memory_init() {
	early_memory_init();
	slab_init();
	init_mapping();
	switch_pt();
}

void* alloc_one_page() {
	u64 i, *bitmap_p;
	for (i = 0; i < mmm.bitmap_num; i++) {
		bitmap_p = mmm.bitmap_p + i;
		if (~(*bitmap_p)) goto found_bitmap;
	}
	return NULL;

found_bitmap:
	for (u64 j = 0; j < 8*sizeof(u64); j++) {
		if (!(*bitmap_p & (UL(1) << j))) {
			*bitmap_p |= (UL(1) << j);
			Page* p = (mmm.pages_p + 8*sizeof(u64)*i + j);
			p->count += 1;
			return p->p_addr;
		}
	}
	return NULL;
}

void free_one_page(void* addr) {
	u64 page_idx = ((u64)addr) >> PAGE_SHIFT;
	u64* bitmap_p = mmm.bitmap_p + (page_idx / (8 * sizeof(u64)));
	*bitmap_p &= ~(UL(1) << (page_idx % (8*sizeof(u64))));
	Page* page_p = mmm.pages_p + page_idx;
	page_p -> count -= 1;
}

void* kmalloc(u64 size) {
	return __kmalloc(size);
}

void kfree(void* addr) {
	__kfree(addr);
}

